#ifndef IP_REASSEMBLY_H
#define IP_REASSEMBLY_H

#include "Packet.h"
#include "LRUList.h"
#include "ip_addr.h"
#include "PointerVector.h"
#include <map>
#include <unordered_set>
#include <memory>
#include "SystemUtils.h"


/**
 * @file
 * This file includes an implementation of IP reassembly mechanism (a.k.a IP de-fragmentation), which is the mechanism of assembling IPv4 or IPv6
 * fragments back into one whole packet. As the previous sentence imply, this module supports both IPv4 and IPv6 reassembly which means
 * the same pcpp#IPReassembly instance can reassemble both IPv4 and IPv6 fragments. You can read more about IP fragmentation here:
 * https://en.wikipedia.org/wiki/IP_fragmentation.<BR>
 * The API is rather simple and contains one main method: pcpp#IPReassembly#processPacket() which gets a fragment packet as a parameter, does the
 * reassembly and returns a fully reassembled packet when done.<BR>
 *
 * The logic works as follows:
 * - There is an internal map that stores the reassembly data for each packet. The key to this map, meaning the way to uniquely associate a
 *   fragment to a (reassembled) packet is the triplet of source IP, destination IP and IP ID (for IPv4) or Fragment ID (for IPv6)
 * - When the first fragment arrives a new record is created in the map and the fragment data is copied
 * - With each fragment arriving the fragment data is copied right after the previous fragment and the reassembled packet is gradually being built
 * - When the last fragment arrives the packet is fully reassembled and returned to the user. Since all fragment data is copied, the packet pointer
 *   returned to the user has to be freed by the user when done using it
 * - The logic supports out-of-order fragments, meaning that a fragment which arrives out-of-order, its data will be copied to a list of out-of-order
 *   fragments where it waits for its turn. This list is observed each time a new fragment arrives to see if the next fragment(s) wait(s) in this
 *   list
 * - If a non-fragment packet arrives it's returned as is to the user
 *
 * In order to limit the amount of memory used by this mechanism there is a limit to the number of concurrent packets being reassembled.
 * The default limit is #PCPP_IP_REASSEMBLY_DEFAULT_MAX_PACKETS_TO_STORE but the user can set any value (determined in pcpp#IPReassembly
 * c'tor). Once capacity (the number of concurrent reassembled packets) exceeds this number, the packet that was least recently used will be
 * dropped from the map along with all the data that was reassembled so far. This means that if the next fragment from this packet suddenly
 * appears it will be treated as a new reassembled packet (which will create another record in the map). The user can be notified when
 * reassembled packets are removed from the map by registering to the pcpp#IPReassembly#OnFragmentsClean callback in pcpp#IPReassembly c'tor
 */

namespace secsmart_tcpip
{    

	#define IP_REASSEMBLY_DEFAULT_MAX_STREAMS_NUM 10000
    #define IP_REASSEMBLY_DEFAULT_MAX_FRAGMENTS_NUM 30000
    #define IP_REASSEMBLY_DEFAULT_TIMEOUT 3000
	
	class IPReassembly
	{
	public:
        struct IPReassemblyCfg
        {
            bool detect_ooo_enable; // 是否检测乱序，默认检测。
            uint32_t max_stream_num; // 分片重组流的规格(默认10000？)
            uint32_t max_frag_cache_num; // 最大缓存分片规格（默认30000？）
            uint32_t timeout; // 老化超时时间,单位ms（默认3秒（3000ms））
        };

        struct IPStreamStatistic
        {
            uint32_t stream_cur_num; // 当前分片流数量
            uint32_t stream_total_num; // 分片流总计
            uint32_t stream_done_num; // 重组完成的分片流
            uint32_t stream_timeout_num; // 超时统计
            uint32_t stream_over_spec_num; // 超规格统计
            uint32_t stream_create_fail; // 创建分片流失败计数
            uint32_t frag_ceate_fail; // 创建缓存分片失败计数
            uint32_t frag_cache_num; // 当前缓存的分片数量
            uint32_t frag_rcv_num; // 分片总计
            uint32_t frag_over_spec_num; // 分片缓存数超规格
            uint32_t frag_ooo_num; // 乱序的分片数量
            uint32_t frag_dup_num; // 重复分片统计
            uint32_t frag_overlap_num; // 重叠分片统计
            uint32_t not_frag_num; // 非分片报文计数
            void clear() { memset(this, 0, sizeof(struct IPStreamStatistic)); }
        };

        struct IPFragStreamId
        {
            uint32_t source; //报文来源 
            uint32_t vlanid; // 在VxLan报文中使用VNI，其他报文使用vlanID。隔离时需要。
            IPAddress sip;
            IPAddress dip;
            uint32_t fragid;

            bool operator<(const IPFragStreamId& id2) const 
            {
                if (source != id2.source) {
                    return source < id2.source;
                }
                if (vlanid != id2.vlanid) {
                    return vlanid < id2.vlanid;
                }
                if (sip != id2.sip) {
                    return sip < id2.sip;
                }
                if (dip != id2.dip) {
                    return dip < id2.dip;
                }
                return fragid < id2.fragid;
            }
        };

        struct IPFragData // 分片数据
		{
            uint8_t *data;
            size_t data_len;
            uint16_t offset;
            bool last;
			IPFragData(uint8_t *data, size_t len,  uint16_t offset, bool last);
			~IPFragData();
		};

		struct  IPFragStream // IP分片流
		{
			IPFragStreamId id;
            uint16_t expect_offset;
			size_t rcv_size;
			size_t total_size;
			bool rcv_last;
			bool rcv_first;
			std::list<IPFragData*> frag_list;
			uint64_t update_time;

			IPFragStream(IPFragStreamId id);
			~IPFragStream();
		};

		enum
		{
            REASSEMBLY_OK = 0, // 重组完成
            REASSEMBLY_NOT_FRAG, // 非分片
            REASSEMBLY_CACHE_FRAG, // 缓存分片
            REASSEMBLY_DROP_FRAG, // 丢弃分片 （重复分片的情况）
            REASSEMBLY_ERROR, // 重组失败（异常报文、规格满了、申请内存失败）
		};

        /**
         * 构造函数
         * @param[in] maxIPStreams：支持IP分片流的规格
         * @param[in] maxFragments：分片缓存数量的规格
         */
		IPReassembly(
            size_t streams_num = IP_REASSEMBLY_DEFAULT_MAX_STREAMS_NUM,
		    size_t frags_num = IP_REASSEMBLY_DEFAULT_MAX_FRAGMENTS_NUM,
		    uint32_t timeout = IP_REASSEMBLY_DEFAULT_TIMEOUT );

		~IPReassembly();

		Packet* process_packet(Packet* fragment, uint32_t &status, ProtocolType parseUntil = UnknownProtocol);
		
        void print_ip_frag_streams();
		
        void set_cfg(uint32_t ip_stream_size, uint32_t frag_cache_size, uint32_t timeout, bool detect_ooo);

        IPReassemblyCfg *get_cfg() { return &m_cfg; }

        IPStreamStatistic *get_statistic() { return &m_stat; }

        uint32_t get_cur_stream_num() { return m_stat.stream_cur_num; }

        uint32_t get_done_stream_num() { return m_stat.stream_done_num; }

        uint32_t get_total_stream_num() { return m_stat.stream_total_num; }

    private:
        typedef std::pair<std::list<IPFragStream*>::iterator, IPFragStream*> IPStreamKeyValuePair;
		std::map<IPFragStreamId, IPStreamKeyValuePair> m_IPFragStreamMap; //存储分片链
		std::list<IPFragStream*> m_AgingList; // 老化链 
		IPReassemblyCfg m_cfg;
        IPStreamStatistic m_stat;
        uint64_t m_last_aging_time; // 上次老化处理时间
        uint32_t m_aging_count_once; // 一次老化处理的流数量
		
        void deleteOldestIPFragStream();
        void updateIPHead(uint8_t version, uint8_t *ip_hdr, size_t hdr_len, size_t data_len, uint8_t proto);
		uint32_t addIPFrag(IPFragStream* fragStream, IPFragData *frag);
        bool detectPktOutOfOrder(IPFragStream* fragStream, IPFragData *frag, std::list<IPFragData*>::iterator pos);
		bool is_complete(IPFragStream* fragStream);
		Packet* reassemblyIPFragStream(IPFragStream* fragStream, Packet* pkt, packet_meta_t *meta);
        void getIPFragStreamId(packet_meta_t *meta, IPFragStreamId *id);
        void aging_proc();
        IPFragStream* getIPFragStream(IPFragStreamId &id);
        IPFragStream* addIPFragStream(IPFragStreamId id);
        void deleteIPFragStream(IPFragStreamId &id);
	};

} // namespace secsmart_tcpip

#endif // PACKETPP_IP_REASSEMBLY
